import asyncio
import cProfile
from datetime import datetime
import time
import logging
from logging.handlers import RotatingFileHandler
import os
import re
import random
import shutil
import threading
from wxauto import WeChat
from openai import OpenAI

# 配置参数集中管理
class Config:
    # 格式：[[账号标识, 群聊名称, 提示词名称], ...]
    LISTEN_LIST = [
        ["1", "AI测试", "提示词示例2"],
        ["2", "AI测试1", "提示词示例2"],
    ]

    DEEPSEEK_API_KEY = "330f54b7-915b-48a9-9599-c5e47c118cb5"  # 替换为实际APIKey
    DEEPSEEK_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
    MODEL = "doubao-1-5-pro-256k-250115"
    MAX_TOKEN = 2000
    TEMPERATURE = 1.1
    QUEUE_WAITING_TIME = 7
    AUTO_MESSAGE = "请你模拟系统设置的角色，在微信上找对方继续刚刚的话题或者询问对方在做什么"
    ENABLE_AUTO_MESSAGE = False
    MIN_COUNTDOWN_HOURS = 1.0
    MAX_COUNTDOWN_HOURS = 2.0
    QUIET_TIME_START = "7:00"
    QUIET_TIME_END = "08:00"
    AVERAGE_TYPING_SPEED = 0.01
    RANDOM_TYPING_SPEED_MIN = 0.01
    RANDOM_TYPING_SPEED_MAX = 0.01
    ENABLE_MEMORY = True
    MEMORY_TEMP_DIR = "Memory_Temp"
    MAX_MESSAGE_LOG_ENTRIES = 4
    MAX_MEMORY_NUMBER = 10
    Accept_All_Group_Chat_Messages = False

# 日志配置优化
log_file = "bot.log"
max_bytes = 10 * 1024 * 1024  # 10MB
backup_count = 5
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        RotatingFileHandler(log_file, maxBytes=max_bytes, backupCount=backup_count),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 多账号微信实例管理（支持多开）
class WeChatAccount:
    def __init__(self, account_id):
        self.account_id = account_id
        self.wx = WeChat(hwnd=int(account_id))  # 每个账号独立微信实例
        self.ROBOT_WX_NAME = self.wx.A_MyIcon.Name
        self._initialize_listen()

    def _initialize_listen(self):
        """初始化监听列表"""
        for _, chat_name, _ in Config.LISTEN_LIST:
            if self.account_id == int(self.get_account_id_from_list(chat_name)):
                try:
                    self.wx.AddListenChat(who=chat_name, savepic=True)
                    logger.info(f"账号{self.account_id}成功监听群聊：{chat_name}")
                except Exception as e:
                    logger.error(f"账号{self.account_id}监听{chat_name}失败: {e}")

    def get_account_id_from_list(self, chat_name):
        """从配置中获取账号ID"""
        for item in Config.LISTEN_LIST:
            if item[1] == chat_name:
                return item[0]
        return None

# 微信管理器（支持多账号）
class WeChatManager:
    _instances = {}
    _lock = threading.Lock()

    @classmethod
    def get_account(cls, account_id):
        """获取或创建微信账号实例"""
        with cls._lock:
            if account_id not in cls._instances:
                cls._instances[account_id] = WeChatAccount(account_id)
            return cls._instances[account_id]

# 用户状态管理（带账号标识）
class UserState:
    def __init__(self):
        self._states = {}
        self._lock = threading.RLock()

    def _get_user_key(self, account_id, chat_name):
        """生成唯一用户标识"""
        return f"{account_id}_{chat_name}"

    def get_state(self, account_id, chat_name):
        """获取用户状态"""
        user_key = self._get_user_key(account_id, chat_name)
        with self._lock:
            return self._states.setdefault(user_key, {
                "messages": [],
                "last_response_time": 0,
                "auto_message_timer": None,
                "memory": []
            })

    def update_timer(self, account_id, chat_name):
        """更新用户计时"""
        user_key = self._get_user_key(account_id, chat_name)
        with self._lock:
            self._states[user_key]["last_response_time"] = time.time()

# 核心功能模块
class ChatBot:
    def __init__(self):
        self.user_state = UserState()
        self.client = OpenAI(
            api_key=Config.DEEPSEEK_API_KEY,
            base_url=Config.DEEPSEEK_BASE_URL
        )
        self._load_config()
        self._init_memory_dir()

    def _load_config(self):
        """加载配置"""
        self.quiet_start = datetime.strptime(Config.QUIET_TIME_START, "%H:%M").time()
        self.quiet_end = datetime.strptime(Config.QUIET_TIME_END, "%H:%M").time()
        self.root_dir = os.path.dirname(os.path.abspath(__file__))
        self.memory_dir = os.path.join(self.root_dir, Config.MEMORY_TEMP_DIR)
        os.makedirs(self.memory_dir, exist_ok=True)

    def _init_memory_dir(self):
        """初始化记忆目录"""
        if not os.path.exists(self.memory_dir):
            os.makedirs(self.memory_dir)

    def is_quiet_time(self):
        """判断免打扰时间"""
        now = datetime.now().time()
        if self.quiet_start <= self.quiet_end:
            return self.quiet_start <= now <= self.quiet_end
        return now >= self.quiet_start or now <= self.quiet_end

    async def process_message(self, msg, account_id, chat_name):
        """处理消息（带账号和群聊标识）"""
        user_state = self.user_state.get_state(account_id, chat_name)
        content = self._clean_message(msg, account_id)
        if not content:
            return

        logger.info(f"[{account_id}/{chat_name}] 收到消息: {content}")
        self._log_memory(account_id, chat_name, f"用户: {content}")
        user_state["messages"].append(content)

        if len(user_state["messages"]) > 10:
            user_state["messages"].pop(0)

        await self._generate_response(account_id, chat_name)

    def _clean_message(self, msg, account_id):
        """清理消息内容"""
        content = msg.content if hasattr(msg, "content") else ""
        if msg.type == "group" and not Config.Accept_All_Group_Chat_Messages:
            robot_name = WeChatManager.get_account(account_id).ROBOT_WX_NAME
            if f"@{robot_name}" in content:
                content = re.sub(f"@{robot_name}", "", content).strip()
                return f"[群聊@{msg.sender}] {content}"
            return None
        return content

    async def _generate_response(self, account_id, chat_name):
        """生成回复"""
        user_state = self.user_state.get_state(account_id, chat_name)
        prompt = self._build_prompt(account_id, chat_name, user_state["messages"])

        try:
            response = self.client.chat.completions.create(
                model=Config.MODEL,
                messages=prompt,
                temperature=Config.TEMPERATURE,
                max_tokens=Config.MAX_TOKEN
            )
            reply = response.choices[0].message.content.strip()
            await self._send_message(account_id, chat_name, reply)
            self._log_memory(account_id, chat_name, f"机器人: {reply}")
            user_state["messages"].append(reply)
        except Exception as e:
            logger.error(f"生成回复失败: {e}")
            await self._send_message(account_id, chat_name, "抱歉，暂时无法回复")

    async def _send_message(self, account_id, chat_name, content):
        """发送消息"""
        if self.is_quiet_time():
            logger.info("免打扰时间，不发送消息")
            return

        try:
            account = WeChatManager.get_account(account_id)
            account.wx.ChatWith(chat_name)
            typing_speed = random.uniform(
                Config.RANDOM_TYPING_SPEED_MIN,
                Config.RANDOM_TYPING_SPEED_MAX
            )
            for part in content.split("\n"):
                if part:
                    account.wx.SendMsg(part,chat_name)
                    await asyncio.sleep(len(part) * typing_speed)
            logger.info(f"[{account_id}/{chat_name}] 消息发送成功")
        except Exception as e:
            logger.error(f"发送消息失败: {e}")

    def _build_prompt(self, account_id, chat_name, messages):
        """构建提示词"""
        prompt_name = next(
            (item[2] for item in Config.LISTEN_LIST
             if item[0] == account_id and item[1] == chat_name),
            "提示词示例2"
        )
        prompt_path = os.path.join(self.root_dir, "prompts", f"{prompt_name}.md")

        try:
            with open(prompt_path, "r", encoding="utf-8") as f:
                system_prompt = f.read()
                return [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": "\n".join(messages)}
                ]
        except FileNotFoundError:
            logger.warning(f"提示词文件{prompt_path}未找到，使用默认提示")
            return [{"role": "user", "content": "\n".join(messages)}]

    def _log_memory(self, account_id, chat_name, content):
        """记录记忆（带账号和群聊标识）"""
        if not Config.ENABLE_MEMORY:
            return

        user_key = f"{account_id}_{chat_name}"
        log_file = os.path.join(self.memory_dir, f"{user_key}_chatlog.txt")

        with open(log_file, "a", encoding="utf-8") as f:
            f.write(f"[{datetime.now()}] {content}\n")

        if os.path.getsize(log_file) > 1024 * 1024:
            shutil.move(log_file, f"{log_file}.old")

    async def start_listening(self):
        """启动多账号监听"""
        account_ids = {item[0] for item in Config.LISTEN_LIST}
        tasks = []

        for account_id in account_ids:
            account = WeChatManager.get_account(int(account_id))
            tasks.append(self._listen_account(account, account_id))

        await asyncio.gather(*tasks)

    async def _listen_account(self, account, account_id):
        """单个账号监听循环"""
        while True:
            try:
                msgs = account.wx.GetListenMessage()
                for chat in msgs:
                    chat_name = chat.who
                    for msg in msgs[chat]:
                        if msg.type in ["friend", "group"]:
                            await self.process_message(msg, account_id, chat_name)
            except Exception as e:
                logger.error(f"账号{account_id}监听异常: {e}")
                await asyncio.sleep(1)
            await asyncio.sleep(0.1)

# 程序入口
async def main():
    try:
        bot = ChatBot()
        logger.info("多账号微信机器人启动")
        await bot.start_listening()
    except KeyboardInterrupt:
        logger.info("用户终止程序")
    except Exception as e:
        logger.error(f"程序异常: {e}", exc_info=True)

if __name__ == "__main__":
    asyncio.run(main())
    # cProfile.run("asyncio.run(main())")